#!/usr/bin/env python
# -*- coding: utf8 -*-
#-------------------------------------------------------------------------------
# Name:        task_from_yied.py
# Purpose: Implement task by python yield
#
# Author:      iori
#
# Created:     30/10/2014
# Copyright:   (c) iori 2014
# Licence:     GPLv3
#-------------------------------------------------------------------------------


from concurrent.futures import ThreadPoolExecutor,  Future, ProcessPoolExecutor
import time
import unittest
from contextlib import contextmanager

@contextmanager
def timeit(label):
    start_time = time.time()
    try:
        yield
    finally:
        print(label, time.time()-start_time)

pool = ThreadPoolExecutor(2)

def patch_future(cls):
    ''' Make Future iterable '''
    def __iter__(self):
        if not self.done():
            yield self
        return self.result()
    cls.__iter__ = __iter__

patch_future(Future) # Make future iterable

class Task(Future):
    def __init__(self, gen):
        super().__init__()
        self._gen = gen

    def step(self, val=None, exc=None):
        try:
            if exc :
                fu = self._gen.throw(exc)
            else:
                fu = self._gen.send(val)
            fu.add_done_callback(self._wakeup)
        # we should no propogate the exception, otherwise it will propogate
        # to `_wakeup` and opps.
        except StopIteration as e: # task end will throw this exception
            self.set_result(e.value) # set future value
        except Exception as e:
            self.set_exception(e) # task abnormal exit.

    def _wakeup(self, fu):
        try:
            res = fu.result()
        except Exception as e:
            self.step(None, exc=e) # propogate exception
        else:
            self.step(res)


def start_inline_task(gen):
    ''' 启动一个异步调用'''
    t = Task(gen)
    t.step()
    return t

def run_inline_task(gen):
    ''' 获取一个异步调用的结果'''
    t = start_inline_task(gen)
    return t.result()


def fib(n):
    ''' 耗时的一个函数 '''
    return 1 if n<3 else (fib(n-1) + fib(n-2))

def fun(f, *nkw):
    ''' 异步执行一个函数 '''
    yield from pool.submit(f, *nkw)



class TestMe(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        pool = ThreadPoolExecutor(2)

    def setUp(self):
        pass

    def test_single(self):
        def fun(x, y):
            yield from pool.submit(time.sleep, 0.1)
            return x + y
        t = Task(fun(3, 6))
        t.step()
        self.assertAlmostEqual(t.result(), 9)

    def test_exception(self):
        def fun(x, y):
            yield from pool.submit(time.sleep, 0.1)
            return x + y
        with self.assertRaises(TypeError):
            t = Task(fun(3, '23'))
            t.step()
            t.result() # wait task to finish

    def test_thread(self):
        def fib(n):
            return 1 if n<3 else (fib(n-1) + fib(n-2))

        def fun(f, n):
            yield from pool.submit(f, n)

        with timeit('thread sequence cost: '):
            run_inline_task(fun(fib, 31))
            run_inline_task(fun(fib, 31))

        with timeit('thread paralell cost: '):
            t1 = start_inline_task(fun(fib, 31))
            t2 = start_inline_task(fun(fib, 31))
            t1.result()
            t2.result()

    def test_process(self):
        ''' 要执行的代码不要放此处定义，'''
        global pool
        thread_pool = pool
        pool = ProcessPoolExecutor(2)
        with timeit('process sequence cost: '):
            run_inline_task(fun(fib, 31))
            run_inline_task(fun(fib, 31))

        with timeit('process paralell cost: '):
            t1 = start_inline_task(fun(fib, 31))
            t2 = start_inline_task(fun(fib, 31))
            t1.result()
            t2.result()
        pool = thread_pool

    def test_timeit(self):
            def fib(n):
                return 1 if n<3 else (fib(n-1) + fib(n-2))

            with timeit('fib: '):
                fib(30)

if __name__ == '__main__':

    unittest.main()




